// SPDX-License-Identifier: GPL-2.0 OR Linux-OpenIB

#include <linux/hash.h>
#include <linux/profile.h>
#include <linux/workqueue.h>
#include <linux/slab.h>

#include "rdma_core.h"
#include "uverbs_profile.h"
#include <trace/hooks/rdma_pangea.h>

struct ufile_hashnode {
	pid_t tgid;
	struct ib_uverbs_file *file;
	struct hlist_node node;
};

struct task_exit_work {
	struct work_struct work;
	struct task_struct *task;
	unsigned long start_time;
};

struct destroy_ufile_work {
	struct work_struct work;
	struct ib_uverbs_file *ufile;
	unsigned long start_time;
};

#define UFILE_HASH_BITS (8)

static DEFINE_HASHTABLE(ufile_hashtable, UFILE_HASH_BITS);
static DECLARE_RWSEM(hash_rwsem);
static struct workqueue_struct *uverbs_profile_wq;
static struct workqueue_struct *uverbs_exit_wq;

#define MONITOR_OS_SCHEDULE(start) \
	do { \
		if (time_after(jiffies, (start) + 1)) { \
			ofed_log_print(PANGEA_WARN, "OS schedule too long, elapsed jiffies %lu.", \
				jiffies - (start)); \
		} \
	} while (0)

int uverbs_profile_add_ufile(struct ib_uverbs_file *ufile)
{
	struct ufile_hashnode *file_node;

	file_node = kzalloc(sizeof(*file_node), GFP_KERNEL);
	if (!file_node) {
		ofed_log_print(PANGEA_ERR, "No memory, kzalloc size(%zu).", sizeof(*file_node));
		return -ENOMEM;
	}

	kref_get(&ufile->ref);
	file_node->tgid = ufile->tgid;
	file_node->file = ufile;
	down_write(&hash_rwsem);
	hlist_add_head(&file_node->node,
		&ufile_hashtable[hash_32(file_node->tgid, UFILE_HASH_BITS)]);
	up_write(&hash_rwsem);

	return 0;
}

void uverbs_profile_remove_ufile(struct ib_uverbs_file *ufile)
{
	struct ufile_hashnode *cur;
	struct hlist_node *tmp;
	struct hlist_head *head;
	pid_t tgid = ufile->tgid;

	down_write(&hash_rwsem);
	head = &ufile_hashtable[hash_32(tgid, UFILE_HASH_BITS)];
	hlist_for_each_entry_safe(cur, tmp, head, node) {
		if (cur->file != ufile)
			continue;

		hlist_del(&cur->node);
		kref_put(&cur->file->ref, ib_uverbs_release_file);
		kfree(cur);
		break;
	}
	up_write(&hash_rwsem);
}

static void destroy_ufile_work_handler(struct work_struct *_work)
{
	struct destroy_ufile_work *work =
		container_of(_work, struct destroy_ufile_work, work);
	struct ib_uverbs_file *ufile = work->ufile;
	unsigned long last_tick = jiffies;

	MONITOR_OS_SCHEDULE(work->start_time);

	// When a CPU is suspended, the SRCU cannot pass the grace period and uses atomic instead.
	synchronize_srcu(&ufile->device->disassociate_srcu);

	MONITOR_OS_SCHEDULE(last_tick);

	ofed_log_print(PANGEA_INFO, "Uverbs profile release ufile(%px).", ufile);
	__uverbs_destroy_ufile_hw(ufile, RDMA_REMOVE_CLOSE, true);
	ofed_log_print(PANGEA_INFO, "Uverbs profile release ufile(%px) done.", ufile);

	kref_put(&ufile->ref, ib_uverbs_release_file);

	kfree(work);
}

static void task_exit_work_handler(struct work_struct *_work)
{
	HLIST_HEAD(tmp_list);
	struct ufile_hashnode *cur;
	struct hlist_node *tmp;
	struct hlist_head *head;
	struct destroy_ufile_work *destroy_work;

	struct task_exit_work *work =
		container_of(_work, struct task_exit_work, work);
	pid_t tgid = work->task->tgid;

	MONITOR_OS_SCHEDULE(work->start_time);

	down_write(&hash_rwsem);
	head = &ufile_hashtable[hash_32(tgid, UFILE_HASH_BITS)];
	hlist_for_each_entry_safe(cur, tmp, head, node) {
		if (cur->tgid != tgid)
			continue;

		hlist_del_init(&cur->node);
		hlist_add_head(&cur->node, &tmp_list);
	}
	up_write(&hash_rwsem);

	head = &tmp_list;
	hlist_for_each_entry_safe(cur, tmp, head, node) {
		// Let other user mode threads return quickly.
		set_file_destroying(cur->file);

		destroy_work = kmalloc(sizeof(*destroy_work), GFP_KERNEL);
		if (!destroy_work) {
			ofed_log_print(PANGEA_ERR, "Can't alloc work for uverbs destroy ufile");
			kref_put(&cur->file->ref, ib_uverbs_release_file);
			goto next;
		}

		INIT_WORK(&destroy_work->work, destroy_ufile_work_handler);
		destroy_work->ufile = cur->file;
		destroy_work->start_time = jiffies;
		queue_work(uverbs_exit_wq, &destroy_work->work);
		ofed_log_print(PANGEA_INFO, "Queue release file(%px) work success.", cur->file);

next:
		hlist_del(&cur->node);
		kfree(cur);
	}

	put_task_struct(work->task);
	kfree(work);
}

static int uverbs_task_exit(struct notifier_block *nb, unsigned long action,
	void *data)
{
	struct task_exit_work *work;
	struct ufile_hashnode *ufile;
	struct hlist_head *head;
	struct task_struct *task = NULL;

	/* Filters out kernel threads */
	if (current->flags & PF_KTHREAD)
		return NOTIFY_OK;

	/* Operates only on the main thread of the process */
	if (current->pid != current->tgid)
		return NOTIFY_OK;

	down_read(&hash_rwsem);
	head = &ufile_hashtable[hash_32(current->tgid, UFILE_HASH_BITS)];
	hlist_for_each_entry(ufile, head, node) {
		if (ufile->tgid == current->tgid) {
			ofed_log_print(PANGEA_INFO, "Found process(%d) exit.", current->tgid);
			get_task_struct(current);
			task = current;
			break;
		}
	}
	up_read(&hash_rwsem);

	if (!task)
		return NOTIFY_OK;

	work = kmalloc(sizeof(*work), GFP_ATOMIC);
	if (!work) {
		ofed_log_print(PANGEA_ERR, "Couldn't allocate work for task exit event, pid(%d).",
			task->tgid);
		put_task_struct(task);
		return NOTIFY_DONE;
	}

	INIT_WORK(&work->work, task_exit_work_handler);
	work->task = task;
	work->start_time = jiffies;

	queue_work(uverbs_profile_wq, &work->work);
	ofed_log_print(PANGEA_INFO, "Uverbs queue work success, process(%d) exit.", current->tgid);

	return NOTIFY_OK;
}

static struct notifier_block uverbs_profile_cb = {
	.notifier_call = uverbs_task_exit,
};

void uverbs_profile_cleanup(void)
{
	profile_event_unregister(PROFILE_TASK_EXIT, &uverbs_profile_cb);
	destroy_workqueue(uverbs_exit_wq);
	destroy_workqueue(uverbs_profile_wq);
}

int uverbs_profile_init(void)
{
	int ret;

	uverbs_profile_wq = alloc_ordered_workqueue("uverbs_profile_wq", 0);
	if (!uverbs_profile_wq) {
		ofed_log_print(PANGEA_ERR, "Alloc workqueue failed.");
		return -ENOMEM;
	}

	uverbs_exit_wq = alloc_workqueue("uverbs_exit_wq",
		WQ_UNBOUND | WQ_HIGHPRI | WQ_MEM_RECLAIM | WQ_SYSFS, 0);
	if (!uverbs_exit_wq) {
		ofed_log_print(PANGEA_ERR, "Alloc uverbs exit wq failed.");
		ret = -ENOMEM;
		goto err;
	}

	ret = profile_event_register(PROFILE_TASK_EXIT, &uverbs_profile_cb);
	if (ret) {
		ofed_log_print(PANGEA_ERR, "Uverbs register profile failed, err(%d).", ret);
		goto err_event;
	}

	return 0;
err_event:
	destroy_workqueue(uverbs_exit_wq);
err:
	destroy_workqueue(uverbs_profile_wq);
	return ret;
}

